package leetcode_2601_2700;

import java.util.*;

public class LeeCode_2607 {
    public static void main(String[] args) {
        System.out.println(makeSubKSumEqual(new int[]{1, 4, 1, 3}, 2));
    }
    private static long makeSubKSumEqual(int[] arr, int k) {
        int n = arr.length;
        int[] p = new int[n];
        for (int i = 0; i < n; i++) {
            p[i] = i;
        }
        int i = 0;
        while (i < n){
            union(p, i, (i + k) % n);
            i++;
        }
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int j = 0; j < n; j++) {
            int root = find(p, j);
            if (map.containsKey(root)){
                map.get(root).add(arr[j]);
            }else {
                List<Integer> list = new ArrayList<>();
                list.add(arr[j]);
                map.put(root, list);
            }
        }
        long ans = 0;
        for (Integer key : map.keySet()) {
            List<Integer> list = map.get(key);
            list.sort(Comparator.comparingInt(o -> o));
            int l = 0, r = list.size() - 1;
            while (l < r){
                ans += list.get(r--) - list.get(l++);
            }
        }
        return ans;
    }
    private static int find(int[] p , int x){
        if (p[x] == x)
            return x;
        return find(p, p[x]);
    }
    private static void union(int[] p, int x, int y){
        int xr = find(p, x);
        int yr = find(p, y);
        p[xr] = yr;
    }
}
